package com.tyrone.algorithm.tree;

public class RedBlackTree<Key extends Comparable<Key>,Value>{
    private Node root ; //记录根结点
    private int N; //记录树中元素的个数
    private static final boolean RED = true; //红色链接标识
    private static final boolean BLACK =false; //黑色链接标识

    /**
     * 判断当前结点的父指向链接是否为红色
     * @param x
     * @return
     */
    private boolean isRed(Node x){
        if (x==null){
            return false;
        }
        return x.color=RED;
    }

    /**
     * 左旋调整
     * @param h
     * @return
     */
    private Node rotateLeft(Node h){
        Node x = h.right;
        Node left = x.left;
        x.left=h;
        h.right=left;
        x.color=h.color;
        h.color=RED;
        return x;
    }

    /**
     * 右旋调整
     * @param h
     * @return
     */
    private Node rotateRight(Node h){
        Node x = h.left;
        Node xRight = x.right;
        x.right=h;
        h.left=xRight;
        x.color=h.color;
        h.color=RED;
        return x;
    }

    /**
     * 颜色反转,相当于完成拆分4-结点
     * @param h
     */
    private void flipColors(Node h){
        h.color=RED;
        h.left.color=BLACK;
        h.right.color=BLACK;
    }

    /**
     * 在整个树上完成插入操作
     * @param key
     * @param val
     */
    public void put(Key key, Value val){
        root=put(root,key,val);
        root.color=BLACK;
    }

    /**
     * 在指定树中，完成插入操作,并返回添加元素后新的树
     * @param h
     * @param key
     * @param val
     * @return
     */
    private Node put(Node h, Key key, Value val){
        if (h==null){
            N++;
            return new Node(key,val,null,null,RED);
        }
        int cmp= key.compareTo(h.key);
        if (cmp<0){
            h.left = put(h.left, key, val);
        }else if(cmp>0){
            h.right=put(h.right,key,val);
        }else {
            h.value=val;
        }
        if (isRed(h.right)&&!isRed(h.left)){
            h=rotateLeft(h);
        }
        if (isRed(h.left)&&isRed(h.left.left)){
            h=rotateRight(h);
        }
        if (isRed(h.left)&&isRed(h.right)){
            flipColors(h);
        }
        return h;
    }

    /**
     * 根据key，从树中找出对应的值
     * @param key
     * @return
     */
    public Value get(Key key){
        return get(root,key);
    }

    /**
     * 从指定的树x中，找出key对应的值
     * @param x
     * @param key
     * @return
     */
    private Value get(Node x, Key key){
        if (x==null){
            return null;
        }
        int cmp = key.compareTo(x.key);
        if (cmp<0){
            return get(x.left,key);
        }else if(cmp>0){
            return get(x.right,key);
        }else {
            return x.value;
        }
    }

    /**
     * 获取树中元素的个数
     * @return
     */
    public int size(){
        return N;
    }

    private class Node{
        private Key key;
        private Value value;
        private Node left;
        private Node right;
        private boolean color;

        public Node(Key key, Value value, Node left, Node right, boolean color) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
            this.color = color;
        }
    }
}
